import kit.hash;

/**
 * Generic HashMap implementation, requiring the key to implement the Hashable trait.
 * Relies on == operator for comparisons, so you may need to overwrite it in custom types.
 */
struct Map[K: Hashable, V] {
    static const internalArrayThreshold: Float = 0.7;

    var length: Size = 0;
    var deleted: Size = 0;
    var allocator: Box[Allocator];
    var internalArray: Array[KeyValuePair[K, V]];

    public static function new(allocator: Box[Allocator], capacity: Int = 16): Map[K, V] using implicit allocator {
        var internalArray: Array[KeyValuePair[K, V]] = Array.new(capacity);
        return struct Self {
            allocator,
            internalArray,
        };
    }

    public function put(key: K, value: V): Ptr[V] {
        var ref = this.putRef(key);
        *ref = value;
        return ref;
    }

    public function putRef(key: K): Ptr[V] {
        var alreadyExists = this.exists(key);
        if (((this.length + this.deleted) as Float) / this.internalArray.length) >= Self.internalArrayThreshold {
            this.resize();
        }
        var index = this.findLocation(key);
        if this.internalArray[index].state == Deleted {
            --this.deleted;
        }
        this.internalArray[index] = struct KeyValuePair[K, V] {
            key: key,
            // value: value,
            state: HashCellState.Occupied,
        };
        if !alreadyExists {
            ++this.length;
        }
        return &this.internalArray[index].value;
    }

    public function get(key: K): Option[V] {
        var index = this.findLocation(key);
        if this.internalArray[index].state == HashCellState.Occupied {
            return Some(this.internalArray[index].value);
        }
        return None;
    }

    public function ref(key: K): Option[Ptr[V]] {
        var index = this.findLocation(key);
        if this.internalArray[index].state == HashCellState.Occupied {
            return Some(&this.internalArray[index].value);
        }
        return None;
    }

    public function remove(key: K): Void {
        var index = this.findLocation(key);
        if this.internalArray[index].state == HashCellState.Occupied && this.internalArray[index].key == key {
            this.internalArray[index].state = HashCellState.Deleted;
            --this.length;
            ++this.deleted;
        }
    }

    public function exists(key: K): Bool {
        var index = this.findLocation(key);
        if this.internalArray[index].state == HashCellState.Occupied && this.internalArray[index].key == key {
            return true;
        }
        return false;
    }

    public function keys(): Array[K] {
        var returnArray: Array[K] = Array.new(this.length);
        var j = 0;
        for slot in this.internalArray {
            if slot.state == HashCellState.Occupied {
                returnArray[j++] = slot.key;
            }
        }
        return returnArray;
    }

    public function free(): Void {
        this.internalArray.free();
    }

    public function clear(): Void {
        this.internalArray.clear();
        this.length = 0;
    }

    function findLocation(key: K): Int {
        var hash = key.Hashable.hash() % this.internalArray.length;
        while this.internalArray[hash].state != HashCellState.Empty && this.internalArray[hash].key != key {
            hash = (hash + 1) % this.internalArray.length;
        }
        return hash;
    }

    function resize(): Void using implicit this.allocator {
        var old = this.internalArray;
        this.internalArray = Array.new((this.internalArray.length * 1.5) as Int);
        for i in 0 ... old.length {
            if old[i].state == HashCellState.Occupied {
                this.put(old[i].key, old[i].value);
                --this.length;
            }
        }
        old.free();
    }

    rules {
        ($this[$k] = $v) => $this.put($k, $v);
        ($this[$k]) => $this.get($k).unwrap();

        (for $ident in $this {$e}) => {
            var __length = $this.internalArray.length;
            for __i in 0 ... __length {
                var __slot = $this.internalArray[__i];
                if __slot.state == HashCellState.Occupied {
                    var $ident = __slot.key;
                    {$e}
                }
            }
        }
    }
}

struct KeyValuePair[K: Hashable, V] {
    public var key: K;
    public var value: V;
    public var state: HashCellState = HashCellState.Empty;
}

// abstract WeakMap[K, V]: Map[K, Shared[V]] {
//     public static function new(allocator: Box[Allocator], capacity: Int): Self using implicit allocator {
//         return Map[K, Shared[V]].new(capacity) as WeakMap[K, V];
//     }

//     public function get(key: K): Option[Shared[V]] {
//         match this.get(key) {
//             Some(ptr) => {
//                 if ptr.active() {
//                     return Some(ptr);
//                 } else {
//                     this.remove(key);
//                     return None;
//                 }
//             }
//         }
//     }

//     public function retain(key: K): Option[Shared[V]] {
//         match this.get(key) {
//             Some(ptr) => {
//                 return Some(ptr.ref());
//             }
//             default => {
//                 return None;
//             }
//         }
//     }

//     public function release(key: K): Bool {
//         match this.get(key) {
//             Some(ptr) => {
//                 return ptr.release();
//             }
//             default => {
//                 return false;
//             }
//         }
//     }
// }
